﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Sockets;
using System.Net;
using System.Timers;


namespace MessageHandling
{
   public class ReadMessageAsync : IDisposable
    {
       public ReadMessageAsync(Socket socket, ReadMessageAsyncCallback callback, int timeout)
       {
           Timer t = new Timer(timeout);
           t.Elapsed += new ElapsedEventHandler(t_Elapsed);

           if (socket == null)
               throw new ArgumentNullException("socket");
           if (callback == null)
               throw new ArgumentNullException("callback");
           userSocket = socket;
           userCallback = callback;
           buffer = new byte[4];
           bytesReceived = 0;
           messageSize = -1;
           IAsyncResult result = socket.BeginReceive(buffer, 0, buffer.Length, SocketFlags.None, new AsyncCallback(OnReceive), this);
           t.Start();
           while (!result.IsCompleted && !_timeout)
               System.Threading.Thread.Sleep(500);
           if (result.IsCompleted)
               t.Stop();
           else
           {
               socket.EndReceive(result);
               ReadMessageEventArgs args = new ReadMessageEventArgs(new TimeoutException("Message timeout"));
               t.Stop();
               callback(args);
           }
       }

       void t_Elapsed(object sender, ElapsedEventArgs e)
       {
           _timeout = true;
       }
        public ReadMessageAsync(Socket socket, ReadMessageAsyncCallback callback)
        {
            
         
            if (socket == null)
                throw new ArgumentNullException("socket");
            if (callback == null)
                throw new ArgumentNullException("callback");


            userSocket = socket;
            userCallback = callback;
            buffer = new byte[4];
            bytesReceived = 0;
            messageSize = -1;
            IAsyncResult result = socket.BeginReceive(buffer, 0, buffer.Length, SocketFlags.None, new AsyncCallback(OnReceive), this);

        }
       
        void OnReceive(IAsyncResult ar)
        {

            ReadMessageAsync state = ar.AsyncState as ReadMessageAsync;

            try
            {

                int count = state.userSocket.EndReceive(ar);
                state.bytesReceived += count;
                if (state.messageSize == -1)//we are still reading the size of the data
                {
                    if (count == 0)
                        throw new ProtocolViolationException("The remote peer closed the connection while reading the message size.");
                    if (state.bytesReceived == 4)//we have received the entire message size information
                    {
                        //read the size of the message
                        state.messageSize = BitConverter.ToInt32(state.buffer, 0);
                        if (state.messageSize < 0)
                        {
                            throw new ProtocolViolationException("The remote peer sent a negative message size.");
                        }
                        state.buffer = new Byte[state.messageSize];
                        //reset the bytes received back to zero
                        //because we are now switching to reading the message body
                        state.bytesReceived = 0;
                    }
                    if (state.messageSize != 0)
                    {
                        //we need more data - could be more of the message size information
                        //or it could be the message body.  The only time we won't need to
                        //read more data is if the message size == 0
                        state.userSocket.BeginReceive(state.buffer,
                        state.bytesReceived, //offset where data can be written
                        state.buffer.Length - state.bytesReceived, //how much data can be read into remaining buffer
                        SocketFlags.None, new AsyncCallback(OnReceive), state);
                    }
                    else
                    {
                        //we have received a "keep alive" message, notify the user...
                        ReadMessageEventArgs args = new ReadMessageEventArgs();
                        state.userCallback(args);
                        state.Dispose();
                    }

                }

                else //we are reading the body of the message
                {

                    if (state.bytesReceived == state.messageSize) //we have the entire message
                    {
                        //notify the user
                        state.userCallback(new ReadMessageEventArgs(state.buffer));
                        //free up our reference to the socket, buffer and the callback object.
                        state.Dispose();
                    }
                    else //need more data.
                    {
                        if (count == 0)
                            throw new ProtocolViolationException("The remote peer closed the connection before the entire message was received");

                        //offset where data can be written
                        //how much data can be read into remaining buffer
                        state.userSocket.BeginReceive(state.buffer,state.bytesReceived, state.buffer.Length - state.bytesReceived, SocketFlags.None, new AsyncCallback(OnReceive), state);
                    }
                }
            }
            catch (Exception ex)
            {
                ReadMessageEventArgs args = new ReadMessageEventArgs(ex);
                state.userCallback(args);
                state.Dispose();
            }
        }

        public void Dispose()
        {
            userSocket = null;
            userCallback = null;
            buffer = null;
        }


        public Socket userSocket;
        public ReadMessageAsyncCallback userCallback;
        public int bytesReceived;
        public byte[] buffer;
        public int messageSize;

       private bool _timeout = false;

    }
}
